package com.gs.coursera.algorithms.core.graph;

import com.gs.coursera.algorithms.core.collection.IndexMinPQ;
import com.gs.coursera.algorithms.core.collection.Queue;

public class PrimMST extends MinimumSpaningTree<Integer, Integer> {

	private Edge<Integer, Integer>[] edgeTo;        // edgeTo[v] = shortest edge from tree vertex to non-tree vertex
    private Integer[] distTo;      // distTo[v] = weight of shortest such edge
    private boolean[] marked;     // marked[v] = true if v on tree, false otherwise
    private IndexMinPQ<Integer> pq;
	
	public PrimMST(Graph<Integer, Integer> graph) {
		super(graph);
		edgeTo = new Edge[graph.getVertexCount()];
        distTo = new Integer[graph.getVertexCount()];
        marked = new boolean[graph.getVertexCount()];
        pq = new IndexMinPQ<Integer>(graph.getVertexCount());
        
        for (int v = 0; v < graph.getVertexCount(); v++) 
        	distTo[v] = Integer.MAX_VALUE;
        for (int v = 0; v < graph.getVertexCount(); v++)      // run from each vertex to find
        {    
        	if (!marked[v]) {
        		prim(graph, v);
        	}
        }
	}

	// run Prim's algorithm in graph G, starting from vertex s
    private void prim(Graph<Integer, Integer> G, int s) {
        distTo[s] = 0;
        pq.insert(s, distTo[s]);
        while (!pq.isEmpty()) {
            int v = pq.delMin();
            scan(G, v);
        }
    }
    
 // scan vertex v
    private void scan(Graph G, int v) {
        marked[v] = true;
        Vertex<Integer, Integer> vertex = G.getVertexByNo(v);
        
        for (Edge<Integer, Integer> e : vertex.getAdjacencyList()) {
            int w = e.getEndVertex().getVertexNumber();
            if (marked[w]) 
            	continue;         // v-w is obsolete edge
            if (e.getWeight() < distTo[w]) {
                distTo[w] = e.getWeight();
                edgeTo[w] = e;
                if (pq.contains(w)) 
                	pq.change(w, distTo[w]);
                else 
                	pq.insert(w, distTo[w]);
            }
        }
    }
	
	@Override
	public Iterable<Edge<Integer, Integer>> findMst() {
		Queue<Edge<Integer,Integer>> mst = new Queue<Edge<Integer,Integer>>();
        for (int v = 0; v < edgeTo.length; v++) {
            Edge<Integer,Integer> e = edgeTo[v];
            if (e != null) {
                mst.enqueue(e);
            }
        }
        return mst;
	}

	@Override
	public Integer getTotalWeight() {
		Integer weight = 0;
		for (int v = 0; v < edgeTo.length; v++) {
            Edge<Integer,Integer> e = edgeTo[v];
            if (e != null) {
            	weight += e.getWeight();
            }
        }
		return weight;
	}

}
